"""Audit logging system for API key operations."""

import functools
from typing import Dict, Any, Optional, List, Callable
from datetime import datetime, timedelta
import motor.motor_asyncio
from fastapi import Request, APIRouter, Depends, Query
from pydantic import BaseModel

from ..models.api_key_model import AuditLog
from ..api.key_management import get_db, get_current_user


router = APIRouter(prefix="/api/audit", tags=["Audit"])


class AuditLogResponse(BaseModel):
    """Audit log response model."""
    timestamp: datetime
    user_id: str
    action: str
    resource_id: Optional[str]
    details: Dict[str, Any]
    ip_address: Optional[str]
    result: str
    error_message: Optional[str]


class AuditLogger:
    """Centralized audit logging."""
    
    def __init__(self, mongodb_client: motor.motor_asyncio.AsyncIOMotorClient = None):
        """Initialize audit logger."""
        self.mongodb = mongodb_client
        self._anomaly_detectors = []
        
    async def log(
        self,
        user_id: str,
        action: str,
        resource_id: Optional[str] = None,
        details: Optional[Dict[str, Any]] = None,
        ip_address: Optional[str] = None,
        user_agent: Optional[str] = None,
        result: str = "success",
        error_message: Optional[str] = None
    ):
        """Log an audit event.
        
        Args:
            user_id: User performing the action
            action: Action performed
            resource_id: Resource being acted upon
            details: Additional details
            ip_address: Client IP address
            user_agent: Client user agent
            result: Result of the action
            error_message: Error message if failed
        """
        audit_log = AuditLog(
            user_id=user_id,
            action=action,
            resource_id=resource_id,
            details=details or {},
            ip_address=ip_address,
            user_agent=user_agent,
            result=result,
            error_message=error_message
        )
        
        # Save to database
        if self.mongodb:
            db = self.mongodb.ai_writing
            await db.audit_logs.insert_one(audit_log.to_dict())
            
        # Check for anomalies
        await self._check_anomalies(audit_log)
        
    async def _check_anomalies(self, log: AuditLog):
        """Check for anomalous behavior.
        
        Args:
            log: Audit log entry
        """
        # Check for suspicious patterns
        if log.action in ["delete", "reveal", "share"]:
            # Count recent similar actions
            if self.mongodb:
                db = self.mongodb.ai_writing
                recent_count = await db.audit_logs.count_documents({
                    "user_id": log.user_id,
                    "action": log.action,
                    "timestamp": {"$gte": datetime.utcnow() - timedelta(hours=1)}
                })
                
                if recent_count > 10:
                    # Trigger anomaly alert
                    await self._trigger_anomaly_alert(
                        f"Excessive {log.action} operations",
                        log
                    )
                    
    async def _trigger_anomaly_alert(self, message: str, log: AuditLog):
        """Trigger anomaly alert.
        
        Args:
            message: Alert message
            log: Related audit log
        """
        # TODO: Integrate with alerting service
        print(f"ANOMALY ALERT: {message} - User: {log.user_id}, Action: {log.action}")
        
    def register_anomaly_detector(self, detector: Callable):
        """Register custom anomaly detector.
        
        Args:
            detector: Async function to detect anomalies
        """
        self._anomaly_detectors.append(detector)


# Global logger instance
audit_logger = AuditLogger()


def audit_log(action: str):
    """Decorator for automatic audit logging.
    
    Args:
        action: Action name to log
    """
    def decorator(func):
        @functools.wraps(func)
        async def wrapper(*args, **kwargs):
            # Extract context (assumes FastAPI endpoint)
            request = None
            current_user = None
            resource_id = None
            
            # Find request and user in args/kwargs
            for arg in args:
                if isinstance(arg, Request):
                    request = arg
                elif isinstance(arg, dict) and "id" in arg:
                    current_user = arg
                    
            for key, value in kwargs.items():
                if key == "request" or key == "req":
                    request = value
                elif key == "current_user":
                    current_user = value
                elif key == "key_id":
                    resource_id = value
                    
            # Execute function
            try:
                result = await func(*args, **kwargs)
                
                # Log success
                if current_user:
                    await audit_logger.log(
                        user_id=current_user.get("id", "unknown"),
                        action=action,
                        resource_id=resource_id,
                        ip_address=request.client.host if request else None,
                        user_agent=request.headers.get("User-Agent") if request else None,
                        result="success"
                    )
                    
                return result
                
            except Exception as e:
                # Log failure
                if current_user:
                    await audit_logger.log(
                        user_id=current_user.get("id", "unknown"),
                        action=action,
                        resource_id=resource_id,
                        ip_address=request.client.host if request else None,
                        user_agent=request.headers.get("User-Agent") if request else None,
                        result="failure",
                        error_message=str(e)
                    )
                raise
                
        return wrapper
    return decorator


@router.get("/keys", response_model=List[AuditLogResponse])
async def get_audit_logs(
    user_id: Optional[str] = Query(None, description="Filter by user"),
    action: Optional[str] = Query(None, description="Filter by action"),
    resource_id: Optional[str] = Query(None, description="Filter by resource"),
    start_date: Optional[datetime] = Query(None, description="Start date"),
    end_date: Optional[datetime] = Query(None, description="End date"),
    result: Optional[str] = Query(None, description="Filter by result"),
    skip: int = Query(0, ge=0),
    limit: int = Query(50, ge=1, le=100),
    current_user: Dict = Depends(get_current_user),
    db=Depends(get_db)
):
    """Query audit logs with filters."""
    # Build query
    query = {}
    
    # Admin can see all, others see only their own
    if not current_user.get("is_admin"):
        query["user_id"] = current_user["id"]
    elif user_id:
        query["user_id"] = user_id
        
    if action:
        query["action"] = action
    if resource_id:
        query["resource_id"] = resource_id
    if result:
        query["result"] = result
        
    if start_date or end_date:
        timestamp_query = {}
        if start_date:
            timestamp_query["$gte"] = start_date
        if end_date:
            timestamp_query["$lte"] = end_date
        query["timestamp"] = timestamp_query
        
    # Query database
    cursor = db.audit_logs.find(query).skip(skip).limit(limit).sort("timestamp", -1)
    logs = []
    
    async for doc in cursor:
        log = AuditLog.from_dict(doc)
        logs.append(AuditLogResponse(
            timestamp=log.timestamp,
            user_id=log.user_id,
            action=log.action,
            resource_id=log.resource_id,
            details=log.details,
            ip_address=log.ip_address,
            result=log.result,
            error_message=log.error_message
        ))
        
    return logs


# Archive job for old logs
async def archive_old_logs(mongodb_client: motor.motor_asyncio.AsyncIOMotorClient,
                          days_to_keep: int = 90):
    """Archive old audit logs to cold storage.
    
    Args:
        mongodb_client: MongoDB client
        days_to_keep: Number of days to keep in hot storage
    """
    db = mongodb_client.ai_writing
    cutoff_date = datetime.utcnow() - timedelta(days=days_to_keep)
    
    # Find logs to archive
    logs_to_archive = []
    cursor = db.audit_logs.find({"timestamp": {"$lt": cutoff_date}})
    
    async for doc in cursor:
        logs_to_archive.append(doc)
        
    if logs_to_archive:
        # TODO: Write to cold storage (S3, etc.)
        print(f"Archiving {len(logs_to_archive)} audit logs")
        
        # Delete from hot storage
        await db.audit_logs.delete_many({"timestamp": {"$lt": cutoff_date}})